{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import sympy as sp\n",
    "from sympy import *  # noqa\n",
    "from sympy.solvers.solveset import linsolve\n",
    "\n",
    "sp.init_session()\n",
    "sp.init_printing()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Global variables:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Dimensional parameters\n",
    "dd = 9\n",
    "DD = 12\n",
    "\n",
    "# Speed of light, time, time step\n",
    "c = sp.symbols(r\"c\", real=True, positive=True)\n",
    "t = sp.symbols(r\"t\", real=True)\n",
    "tn = sp.symbols(r\"t_n\", real=True)\n",
    "dt = sp.symbols(r\"\\Delta{t}\", real=True, positive=True)\n",
    "\n",
    "# Components of k vector, omega, norm of k vector\n",
    "kx = sp.symbols(r\"k_x\", real=True)\n",
    "ky = sp.symbols(r\"k_y\", real=True)\n",
    "kz = sp.symbols(r\"k_z\", real=True)\n",
    "om = sp.symbols(r\"omega\", real=True, positive=True)\n",
    "knorm = sp.sqrt(kx**2 + ky**2 + kz**2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Method:\n",
    "The goal is to solve a system of second-order ordinary differential equations in time of the form $\\boldsymbol{\\ddot{X}} = M \\boldsymbol{X}$, where the $d \\times d$ matrix $M$ has eigenpairs $(\\lambda_i, \\{\\boldsymbol{v}_{i,1}, \\dots, \\boldsymbol{v}_{i,\\mu_i}\\})$, where $\\mu_i$ denotes the algebraic multiplicity of $\\lambda_i$ and $\\sum_i \\mu_i = d$.\n",
    "Assuming that all eigenvalues are either zero or negative, we can write $\\lambda_i = - \\omega_i^2$, with $\\omega_i = \\sqrt{- \\lambda_i} \\geq 0$.\n",
    "Then the solution of $\\boldsymbol{\\ddot{X}} = M \\boldsymbol{X}$ reads\n",
    "$$\n",
    "\\boldsymbol{X} = \\sum_i \\sum_j^{\\mu_i} \\left(a_{i,j} C(\\omega_i, t) + b_{i,j} S(\\omega_i, t) \\right) \\boldsymbol{v}_{i,j} \\,,\n",
    "$$\n",
    "where $a_{i,j}$ and $b_{i,j}$ are integration constants to be determined by the initial conditions $\\boldsymbol{X}(t_n)$ and $\\boldsymbol{\\dot{X}}(t_n)$, $C(\\omega, t) = \\cos(\\omega \\, (t - t_n))$, and\n",
    "$$\n",
    "S(\\omega, t) =\n",
    "\\begin{cases}\n",
    "(t - t_n) & \\omega = 0 \\,, \\\\\n",
    "\\dfrac{\\sin(\\omega \\, (t - t_n))}{\\omega} & \\omega \\neq 0 \\,.\n",
    "\\end{cases}\n",
    "$$\n",
    "We remark that\n",
    "$$\n",
    "\\begin{aligned}\n",
    "& \\boldsymbol{X}(t_n) = \\sum_i \\sum_j^{\\mu_i} \\left(a_{i,j} C(\\omega_i, t_n) + b_{i,j} S(\\omega_i, t_n) \\right) \\boldsymbol{v}_{i,j} = \\sum_i \\sum_j^{\\mu_i} a_{i,j} \\boldsymbol{v}_{i,j} \\,, \\\\\n",
    "& \\boldsymbol{\\dot{X}}(t_n) = \\sum_i \\sum_j^{\\mu_i} \\left(a_{i,j} \\dot{C}(\\omega_i, t_n) + b_{i,j} \\dot{S}(\\omega_i, t_n) \\right) \\boldsymbol{v}_{i,j} = \\sum_i \\sum_j^{\\mu_i} b_{i,j} \\boldsymbol{v}_{i,j} \\,.\n",
    "\\end{aligned}\n",
    "$$\n",
    "In fact, the second time derivative of $\\boldsymbol{X}$ yields\n",
    "$$\n",
    "\\begin{split}\n",
    "\\boldsymbol{\\ddot{X}} & = \\sum_i \\sum_j^{\\mu_i} \\left(a_{i,j} \\ddot{C}(\\omega_i, t) + b_{i,j} \\ddot{S}(\\omega_i, t) \\right) \\boldsymbol{v}_{i,j} \\\\\n",
    "& = \\sum_i (-\\omega_i^2) \\sum_j^{\\mu_i} \\left(a_{i,j} C(\\omega_i, t) + b_{i,j} S(\\omega_i, t) \\right) \\boldsymbol{v}_{i,j} = M \\boldsymbol{X} \\,,\n",
    "\\end{split}\n",
    "$$\n",
    "where we used the fact that, by definition, $M \\boldsymbol{v}_{i,j} = \\lambda_i \\boldsymbol{v}_{i,j} = - \\omega_i^2 \\boldsymbol{v}_{i,j}$ for all $j = 1, \\dots, \\mu_i$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Auxiliary functions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def C(omega, t):\n",
    "    return sp.cos(omega * (t - tn))\n",
    "\n",
    "\n",
    "def S(omega, t):\n",
    "    return (t - tn) if omega == 0.0 else sp.sin(omega * (t - tn)) / omega\n",
    "\n",
    "\n",
    "def Xt(eigenpairs, a, b, t):\n",
    "    \"\"\"\n",
    "    Compute X(t) according to the formulas above\n",
    "    for a given set of eigenpairs and coefficients.\n",
    "    \"\"\"\n",
    "    XX = zeros(DD, 1)\n",
    "    # Index used for the integration constants a_n and b_n\n",
    "    i = 0\n",
    "    # Loop over matrix eigenpairs\n",
    "    for ep in eigenpairs:\n",
    "        # ep[0] is an eigenvalue and om = sp.sqrt(-ep[0])\n",
    "        omega = 0.0 if ep[0] == 0.0 else om\n",
    "        # am is the algebraic multiplicity of the eigenvalue\n",
    "        am = ep[1]\n",
    "        # vF is the list of all eigenvectors corresponding to the eigenvalue\n",
    "        vX = ep[2]\n",
    "        # Loop over algebraic multiplicity of the eigenvalue\n",
    "        for j in range(am):\n",
    "            XX += (a[i] * C(omega, t) + b[i] * S(omega, t)) * vX[j]\n",
    "            i += 1\n",
    "    return XX\n",
    "\n",
    "\n",
    "def evolve(X, dX_dt, d2X_dt2):\n",
    "    \"\"\"\n",
    "    Solve ordinary differential equation X'' = M*X.\n",
    "    \"\"\"\n",
    "    # Set matrix for given ODE\n",
    "    MX = zeros(DD)\n",
    "    for i in range(DD):\n",
    "        for j in range(DD):\n",
    "            MX[i, j] = d2X_dt2[i].coeff(X[j], 1)\n",
    "    # MX /= c**2\n",
    "\n",
    "    print()\n",
    "    print(r\"Matrix:\")\n",
    "    display(MX)\n",
    "\n",
    "    # Characteristic matrix\n",
    "    lamda = sp.symbols(r\"lamda\")\n",
    "    Id = eye(DD)\n",
    "    MX_charmat = MX - lamda * Id\n",
    "\n",
    "    # Characteristic polynomial\n",
    "    MX_charpoly = MX_charmat.det()\n",
    "    MX_charpoly = factor(MX_charpoly.as_expr())\n",
    "\n",
    "    print(r\"Characteristic polynomial:\")\n",
    "    display(MX_charpoly)\n",
    "\n",
    "    MX_eigenvals = sp.solve(MX_charpoly, lamda)\n",
    "\n",
    "    # List of eigenvectors\n",
    "    MX_eigenvects = []\n",
    "\n",
    "    # List of eigenpairs\n",
    "    MX_eigenpairs = []\n",
    "\n",
    "    # Compute eigenvectors as null spaces\n",
    "    for ev in MX_eigenvals:\n",
    "        # M - lamda * Id\n",
    "        A = MX_charmat.subs(lamda, ev)\n",
    "        A.simplify()\n",
    "\n",
    "        print(r\"Eigenvalue:\")\n",
    "        display(ev)\n",
    "\n",
    "        print(r\"Characteristic matrix:\")\n",
    "        display(A)\n",
    "\n",
    "        # Perform Gaussian elimination (necessary for lamda != 0)\n",
    "        if ev != 0.0:\n",
    "            print(r\"Gaussian elimination:\")\n",
    "            print(r\"A[0,:] += A[1,:]\")\n",
    "            A[0, :] += A[1, :]\n",
    "            print(r\"A[0,:] += A[2,:]\")\n",
    "            A[0, :] += A[2, :]\n",
    "            print(r\"Swap A[0,:] and A[11,:]\")\n",
    "            row = A[11, :]\n",
    "            A[11, :] = A[0, :]\n",
    "            A[0, :] = row\n",
    "            print(r\"A[3,:] += A[4,:]\")\n",
    "            A[3, :] += A[4, :]\n",
    "            print(r\"A[3,:] += A[5,:]\")\n",
    "            A[3, :] += A[5, :]\n",
    "            print(r\"Swap A[3,:] and A[10,:]\")\n",
    "            row = A[10, :]\n",
    "            A[10, :] = A[3, :]\n",
    "            A[3, :] = row\n",
    "            print(r\"A[0,:] += A[3,:]\")\n",
    "            A[0, :] += A[3, :]\n",
    "            print(r\"A[0,:] += A[9,:]\")\n",
    "            A[0, :] += A[9, :]\n",
    "            print(r\"Swap A[0,:] and A[9,:]\")\n",
    "            row = A[9, :]\n",
    "            A[9, :] = A[0, :]\n",
    "            A[0, :] = row\n",
    "            print(r\"A[6,:] += A[8,:]\")\n",
    "            A[6, :] += A[8, :]\n",
    "            print(r\"A[6,:] += A[7,:]\")\n",
    "            A[6, :] += A[7, :]\n",
    "            print(r\"Swap A[6,:] and A[8,:]\")\n",
    "            row = A[8, :]\n",
    "            A[8, :] = A[6, :]\n",
    "            A[6, :] = row\n",
    "            print(r\"A[6,:] += A[7,:]\")\n",
    "            A[6, :] += A[7, :]\n",
    "            print(r\"A[4,:] += A[5,:]\")\n",
    "            A[4, :] += A[5, :]\n",
    "            A.simplify()\n",
    "            display(A)\n",
    "\n",
    "        # Compute null space and store eigenvectors\n",
    "        v = A.nullspace()\n",
    "        MX_eigenvects.append(v)\n",
    "\n",
    "        print(r\"Eigenvectors:\")\n",
    "        display(v)\n",
    "\n",
    "        # Store eigenpairs (eigenvalue, algebraic multiplicity, eigenvectors)\n",
    "        MX_eigenpairs.append((ev, len(v), v))\n",
    "\n",
    "        # print(r'Eigenpairs:')\n",
    "        # display(MX_eigenpairs)\n",
    "\n",
    "    # Verify that the eigenpairs satisfy the characteristic equations\n",
    "    for ep in MX_eigenpairs:\n",
    "        for j in range(ep[1]):\n",
    "            diff = MX * ep[2][j] - ep[0] * ep[2][j]\n",
    "            diff.simplify()\n",
    "            if diff != zeros(DD, 1):\n",
    "                print(\"The charcteristic equation is not verified for some eigenpairs\")\n",
    "                display(diff)\n",
    "\n",
    "    # Define integration constants\n",
    "    a = []\n",
    "    b = []\n",
    "    for i in range(DD):\n",
    "        an = r\"a_{:d}\".format(i + 1)\n",
    "        bn = r\"b_{:d}\".format(i + 1)\n",
    "        a.append(sp.symbols(an))\n",
    "        b.append(sp.symbols(bn))\n",
    "\n",
    "    # Set equations corresponding to initial conditions\n",
    "    lhs_a = Xt(MX_eigenpairs, a, b, tn) - X\n",
    "    lhs_b = Xt(MX_eigenpairs, a, b, t).diff(t).subs(t, tn) - dX_dt\n",
    "\n",
    "    # Compute integration constants from initial conditions\n",
    "    # (convert list of tuples to list using list comprehension)\n",
    "    a = list(linsolve(list(lhs_a), a))\n",
    "    a = [item for el in a for item in el]\n",
    "    b = list(linsolve(list(lhs_b), b))\n",
    "    b = [item for el in b for item in el]\n",
    "\n",
    "    # Evaluate solution at t = tn + dt\n",
    "    X_new = Xt(MX_eigenpairs, a, b, tn + dt).expand()\n",
    "    for d in range(DD):\n",
    "        for Eij in E:\n",
    "            X_new[d] = X_new[d].collect(Eij)\n",
    "        for Bij in B:\n",
    "            X_new[d] = X_new[d].collect(Bij)\n",
    "        for Fi in F:\n",
    "            X_new[d] = X_new[d].collect(Fi)\n",
    "        for Gi in G:\n",
    "            X_new[d] = X_new[d].collect(Gi)\n",
    "\n",
    "    # Check correctness by taking *second* derivative\n",
    "    # and comparing with initial right-hand side at time tn\n",
    "    X_t = Xt(MX_eigenpairs, a, b, t)\n",
    "    diff = X_t.diff(t).diff(t).subs(t, tn).subs(om, c * knorm).expand() - d2X_dt2\n",
    "    diff.simplify()\n",
    "    if diff != zeros(DD, 1):\n",
    "        print(\"Integration in time failed\")\n",
    "        display(diff)\n",
    "\n",
    "    return X_t, X_new"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### First-order and second-order ODEs for $\\boldsymbol{E}$, $\\boldsymbol{B}$, $F$ and $G$:\n",
    "Equations for the $\\boldsymbol{E}$ field:\n",
    "$$\n",
    "\\begin{alignat*}{3}\n",
    "&  \\frac{\\partial E_{xx}}{\\partial t} = c^2 i k_x (F_x + F_y  + F_z) \\quad\n",
    "&& \\frac{\\partial E_{xy}}{\\partial t} = c^2 i k_y (B_{zx} + B_{zy} + B_{zz}) \\quad\n",
    "&& \\frac{\\partial E_{xz}}{\\partial t} = -c^2 i k_z (B_{yx} + B_{yy} + B_{yz}) \\\\[5pt]\n",
    "&  \\frac{\\partial E_{yx}}{\\partial t} = -c^2 i k_x (B_{zx} + B_{zy} + B_{zz}) \\quad\n",
    "&& \\frac{\\partial E_{yy}}{\\partial t} = c^2 i k_y (F_x + F_y + F_z) \\quad\n",
    "&& \\frac{\\partial E_{yz}}{\\partial t} = c^2 i k_z (B_{xx} + B_{xy} + B_{xz}) \\\\[5pt]\n",
    "&  \\frac{\\partial E_{zx}}{\\partial t} = c^2 i k_x (B_{yx} + B_{yy} + B_{yz}) \\quad\n",
    "&& \\frac{\\partial E_{zy}}{\\partial t} = -c^2 i k_y (B_{xx} + B_{xy} + B_{xz})\\quad\n",
    "&& \\frac{\\partial E_{zz}}{\\partial t} = c^2 i k_z (F_x + F_y + F_z)\n",
    "\\end{alignat*}\n",
    "$$\n",
    "\n",
    "Equations for the $\\boldsymbol{B}$ field:\n",
    "$$\n",
    "\\begin{alignat*}{3}\n",
    "&  \\frac{\\partial B_{xx}}{\\partial t} = i k_x (G_x + G_y + G_z) \\quad\n",
    "&& \\frac{\\partial B_{xy}}{\\partial t} = -i k_y (E_{zx} + E_{zy} + E_{zz}) \\quad\n",
    "&& \\frac{\\partial B_{xz}}{\\partial t} = i k_z (E_{yx} + E_{yy} + E_{yz}) \\\\[5pt]\n",
    "&  \\frac{\\partial B_{yx}}{\\partial t} = i k_x (E_{zx} + E_{zy} + E_{zz}) \\quad\n",
    "&& \\frac{\\partial B_{yy}}{\\partial t} = i k_y (G_x + G_y + G_z) \\quad\n",
    "&& \\frac{\\partial B_{yz}}{\\partial t} = -i k_z (E_{xx} + E_{xy} + E_{xz}) \\\\[5pt]\n",
    "&  \\frac{\\partial B_{zx}}{\\partial t} = -i k_x (E_{yx} + E_{yy} + E_{yz}) \\quad\n",
    "&& \\frac{\\partial B_{zy}}{\\partial t} = i k_y (E_{xx} + E_{xy} + E_{xz}) \\quad\n",
    "&& \\frac{\\partial B_{zz}}{\\partial t} = i k_z (G_x + G_y + G_z)\n",
    "\\end{alignat*}\n",
    "$$\n",
    "\n",
    "Equations for the $F$ field:\n",
    "$$\n",
    "\\frac{\\partial F_x}{\\partial t} = i k_x (E_{xx} + E_{xy} + E_{xz}) \\quad\n",
    "\\frac{\\partial F_y}{\\partial t} = i k_y (E_{yx} + E_{yy} + E_{yz}) \\quad\n",
    "\\frac{\\partial F_z}{\\partial t} = i k_z (E_{zx} + E_{zy} + E_{zz})\n",
    "$$\n",
    "\n",
    "Equations for the $G$ field:\n",
    "$$\n",
    "\\frac{\\partial G_x}{\\partial t} = c^2 i k_x (B_{xx} + B_{xy} + B_{xz}) \\quad\n",
    "\\frac{\\partial G_y}{\\partial t} = c^2 i k_y (B_{yx} + B_{yy} + B_{yz}) \\quad\n",
    "\\frac{\\partial G_z}{\\partial t} = c^2 i k_z (B_{zx} + B_{zy} + B_{zz})\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# indices  0     1     2     3     4     5     6     7     8\n",
    "labels = [\"xx\", \"xy\", \"xz\", \"yx\", \"yy\", \"yz\", \"zx\", \"zy\", \"zz\"]\n",
    "\n",
    "# E fields\n",
    "Exx = sp.symbols(r\"E_{xx}\")\n",
    "Exy = sp.symbols(r\"E_{xy}\")\n",
    "Exz = sp.symbols(r\"E_{xz}\")\n",
    "Eyx = sp.symbols(r\"E_{yx}\")\n",
    "Eyy = sp.symbols(r\"E_{yy}\")\n",
    "Eyz = sp.symbols(r\"E_{yz}\")\n",
    "Ezx = sp.symbols(r\"E_{zx}\")\n",
    "Ezy = sp.symbols(r\"E_{zy}\")\n",
    "Ezz = sp.symbols(r\"E_{zz}\")\n",
    "E = Matrix([[Exx], [Exy], [Exz], [Eyx], [Eyy], [Eyz], [Ezx], [Ezy], [Ezz]])\n",
    "\n",
    "# B fields\n",
    "Bxx = sp.symbols(r\"B_{xx}\")\n",
    "Bxy = sp.symbols(r\"B_{xy}\")\n",
    "Bxz = sp.symbols(r\"B_{xz}\")\n",
    "Byx = sp.symbols(r\"B_{yx}\")\n",
    "Byy = sp.symbols(r\"B_{yy}\")\n",
    "Byz = sp.symbols(r\"B_{yz}\")\n",
    "Bzx = sp.symbols(r\"B_{zx}\")\n",
    "Bzy = sp.symbols(r\"B_{zy}\")\n",
    "Bzz = sp.symbols(r\"B_{zz}\")\n",
    "B = Matrix([[Bxx], [Bxy], [Bxz], [Byx], [Byy], [Byz], [Bzx], [Bzy], [Bzz]])\n",
    "\n",
    "# F fields\n",
    "Fx = sp.symbols(r\"F_{x}\")\n",
    "Fy = sp.symbols(r\"F_{y}\")\n",
    "Fz = sp.symbols(r\"F_{z}\")\n",
    "F = Matrix([[Fx], [Fy], [Fz]])\n",
    "\n",
    "# G fields\n",
    "Gx = sp.symbols(r\"G_{x}\")\n",
    "Gy = sp.symbols(r\"G_{y}\")\n",
    "Gz = sp.symbols(r\"G_{z}\")\n",
    "G = Matrix([[Gx], [Gy], [Gz]])\n",
    "\n",
    "# dE/dt\n",
    "dExx_dt = c**2 * I * kx * (Fx + Fy + Fz)\n",
    "dExy_dt = c**2 * I * ky * (Bzx + Bzy + Bzz)\n",
    "dExz_dt = -(c**2) * I * kz * (Byx + Byy + Byz)\n",
    "dEyx_dt = -(c**2) * I * kx * (Bzx + Bzy + Bzz)\n",
    "dEyy_dt = c**2 * I * ky * (Fx + Fy + Fz)\n",
    "dEyz_dt = c**2 * I * kz * (Bxx + Bxy + Bxz)\n",
    "dEzx_dt = c**2 * I * kx * (Byx + Byy + Byz)\n",
    "dEzy_dt = -(c**2) * I * ky * (Bxx + Bxy + Bxz)\n",
    "dEzz_dt = c**2 * I * kz * (Fx + Fy + Fz)\n",
    "dE_dt = Matrix(\n",
    "    [\n",
    "        [dExx_dt],\n",
    "        [dExy_dt],\n",
    "        [dExz_dt],\n",
    "        [dEyx_dt],\n",
    "        [dEyy_dt],\n",
    "        [dEyz_dt],\n",
    "        [dEzx_dt],\n",
    "        [dEzy_dt],\n",
    "        [dEzz_dt],\n",
    "    ]\n",
    ")\n",
    "\n",
    "# dB/dt\n",
    "dBxx_dt = I * kx * (Gx + Gy + Gz)\n",
    "dBxy_dt = -I * ky * (Ezx + Ezy + Ezz)\n",
    "dBxz_dt = I * kz * (Eyx + Eyy + Eyz)\n",
    "dByx_dt = I * kx * (Ezx + Ezy + Ezz)\n",
    "dByy_dt = I * ky * (Gx + Gy + Gz)\n",
    "dByz_dt = -I * kz * (Exx + Exy + Exz)\n",
    "dBzx_dt = -I * kx * (Eyx + Eyy + Eyz)\n",
    "dBzy_dt = I * ky * (Exx + Exy + Exz)\n",
    "dBzz_dt = I * kz * (Gx + Gy + Gz)\n",
    "dB_dt = Matrix(\n",
    "    [\n",
    "        [dBxx_dt],\n",
    "        [dBxy_dt],\n",
    "        [dBxz_dt],\n",
    "        [dByx_dt],\n",
    "        [dByy_dt],\n",
    "        [dByz_dt],\n",
    "        [dBzx_dt],\n",
    "        [dBzy_dt],\n",
    "        [dBzz_dt],\n",
    "    ]\n",
    ")\n",
    "\n",
    "# dF/dt\n",
    "dFx_dt = I * kx * (Exx + Exy + Exz)\n",
    "dFy_dt = I * ky * (Eyx + Eyy + Eyz)\n",
    "dFz_dt = I * kz * (Ezx + Ezy + Ezz)\n",
    "dF_dt = Matrix([[dFx_dt], [dFy_dt], [dFz_dt]])\n",
    "\n",
    "# dG/dt\n",
    "dGx_dt = c**2 * I * kx * (Bxx + Bxy + Bxz)\n",
    "dGy_dt = c**2 * I * ky * (Byx + Byy + Byz)\n",
    "dGz_dt = c**2 * I * kz * (Bzx + Bzy + Bzz)\n",
    "dG_dt = Matrix([[dGx_dt], [dGy_dt], [dGz_dt]])\n",
    "\n",
    "# d2E/dt2\n",
    "d2Exx_dt2 = c**2 * I * kx * (dFx_dt + dFy_dt + dFz_dt)\n",
    "d2Exy_dt2 = c**2 * I * ky * (dBzx_dt + dBzy_dt + dBzz_dt)\n",
    "d2Exz_dt2 = -(c**2) * I * kz * (dByx_dt + dByy_dt + dByz_dt)\n",
    "d2Eyx_dt2 = -(c**2) * I * kx * (dBzx_dt + dBzy_dt + dBzz_dt)\n",
    "d2Eyy_dt2 = c**2 * I * ky * (dFx_dt + dFy_dt + dFz_dt)\n",
    "d2Eyz_dt2 = c**2 * I * kz * (dBxx_dt + dBxy_dt + dBxz_dt)\n",
    "d2Ezx_dt2 = c**2 * I * kx * (dByx_dt + dByy_dt + dByz_dt)\n",
    "d2Ezy_dt2 = -(c**2) * I * ky * (dBxx_dt + dBxy_dt + dBxz_dt)\n",
    "d2Ezz_dt2 = c**2 * I * kz * (dFx_dt + dFy_dt + dFz_dt)\n",
    "d2E_dt2 = Matrix(\n",
    "    [\n",
    "        [d2Exx_dt2],\n",
    "        [d2Exy_dt2],\n",
    "        [d2Exz_dt2],\n",
    "        [d2Eyx_dt2],\n",
    "        [d2Eyy_dt2],\n",
    "        [d2Eyz_dt2],\n",
    "        [d2Ezx_dt2],\n",
    "        [d2Ezy_dt2],\n",
    "        [d2Ezz_dt2],\n",
    "    ]\n",
    ")\n",
    "\n",
    "# d2B/dt2\n",
    "d2Bxx_dt2 = I * kx * (dGx_dt + dGy_dt + dGz_dt)\n",
    "d2Bxy_dt2 = -I * ky * (dEzx_dt + dEzy_dt + dEzz_dt)\n",
    "d2Bxz_dt2 = I * kz * (dEyx_dt + dEyy_dt + dEyz_dt)\n",
    "d2Byx_dt2 = I * kx * (dEzx_dt + dEzy_dt + dEzz_dt)\n",
    "d2Byy_dt2 = I * ky * (dGx_dt + dGy_dt + dGz_dt)\n",
    "d2Byz_dt2 = -I * kz * (dExx_dt + dExy_dt + dExz_dt)\n",
    "d2Bzx_dt2 = -I * kx * (dEyx_dt + dEyy_dt + dEyz_dt)\n",
    "d2Bzy_dt2 = I * ky * (dExx_dt + dExy_dt + dExz_dt)\n",
    "d2Bzz_dt2 = I * kz * (dGx_dt + dGy_dt + dGz_dt)\n",
    "d2B_dt2 = Matrix(\n",
    "    [\n",
    "        [d2Bxx_dt2],\n",
    "        [d2Bxy_dt2],\n",
    "        [d2Bxz_dt2],\n",
    "        [d2Byx_dt2],\n",
    "        [d2Byy_dt2],\n",
    "        [d2Byz_dt2],\n",
    "        [d2Bzx_dt2],\n",
    "        [d2Bzy_dt2],\n",
    "        [d2Bzz_dt2],\n",
    "    ]\n",
    ")\n",
    "\n",
    "# d2F/dt2\n",
    "d2Fx_dt2 = I * kx * (dExx_dt + dExy_dt + dExz_dt)\n",
    "d2Fy_dt2 = I * ky * (dEyx_dt + dEyy_dt + dEyz_dt)\n",
    "d2Fz_dt2 = I * kz * (dEzx_dt + dEzy_dt + dEzz_dt)\n",
    "d2F_dt2 = Matrix([[d2Fx_dt2], [d2Fy_dt2], [d2Fz_dt2]])\n",
    "\n",
    "# d2G/dt2\n",
    "d2Gx_dt2 = c**2 * I * kx * (dBxx_dt + dBxy_dt + dBxz_dt)\n",
    "d2Gy_dt2 = c**2 * I * ky * (dByx_dt + dByy_dt + dByz_dt)\n",
    "d2Gz_dt2 = c**2 * I * kz * (dBzx_dt + dBzy_dt + dBzz_dt)\n",
    "d2G_dt2 = Matrix([[d2Gx_dt2], [d2Gy_dt2], [d2Gz_dt2]])\n",
    "\n",
    "for i in range(dd):\n",
    "    d2E_dt2[i] = sp.expand(d2E_dt2[i])\n",
    "\n",
    "for i in range(dd):\n",
    "    d2B_dt2[i] = sp.expand(d2B_dt2[i])\n",
    "\n",
    "for i in range(3):\n",
    "    d2F_dt2[i] = sp.expand(d2F_dt2[i])\n",
    "\n",
    "for i in range(3):\n",
    "    d2G_dt2[i] = sp.expand(d2G_dt2[i])\n",
    "\n",
    "# Extended array for E and G\n",
    "EG = zeros(DD, 1)\n",
    "for i in range(dd):\n",
    "    EG[i] = E[i]\n",
    "for i in range(dd, DD):\n",
    "    EG[i] = G[i - dd]\n",
    "\n",
    "# dEG/dt\n",
    "dEG_dt = zeros(DD, 1)\n",
    "for i in range(dd):\n",
    "    dEG_dt[i] = dE_dt[i]\n",
    "for i in range(dd, DD):\n",
    "    dEG_dt[i] = dG_dt[i - dd]\n",
    "\n",
    "# d2EG/dt2\n",
    "d2EG_dt2 = zeros(DD, 1)\n",
    "for i in range(dd):\n",
    "    d2EG_dt2[i] = d2E_dt2[i]\n",
    "for i in range(dd, DD):\n",
    "    d2EG_dt2[i] = d2G_dt2[i - dd]\n",
    "\n",
    "# Extended array for B and F\n",
    "BF = zeros(DD, 1)\n",
    "for i in range(dd):\n",
    "    BF[i] = B[i]\n",
    "for i in range(dd, DD):\n",
    "    BF[i] = F[i - dd]\n",
    "\n",
    "# dBF/dt\n",
    "dBF_dt = zeros(DD, 1)\n",
    "for i in range(dd):\n",
    "    dBF_dt[i] = dB_dt[i]\n",
    "for i in range(dd, DD):\n",
    "    dBF_dt[i] = dF_dt[i - dd]\n",
    "\n",
    "# d2BF/dt2\n",
    "d2BF_dt2 = zeros(DD, 1)\n",
    "for i in range(dd):\n",
    "    d2BF_dt2[i] = d2B_dt2[i]\n",
    "for i in range(dd, DD):\n",
    "    d2BF_dt2[i] = d2F_dt2[i - dd]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Solve second-order ODEs for $\\boldsymbol{E}$, $\\boldsymbol{B}$, $F$ and $G$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(r\"Solve equations for E and G:\")\n",
    "EG_t, EG_new = evolve(EG, dEG_dt, d2EG_dt2)\n",
    "\n",
    "print(r\"Solve equations for B and F:\")\n",
    "BF_t, BF_new = evolve(BF, dBF_dt, d2BF_dt2)\n",
    "\n",
    "# Check correctness by taking *first* derivative\n",
    "# and comparing with initial right-hand side at time tn\n",
    "# E,G\n",
    "diff = EG_t.diff(t).subs(t, tn).subs(om, c * knorm).expand() - dEG_dt\n",
    "diff.simplify()\n",
    "if diff != zeros(DD, 1):\n",
    "    print(\"Integration in time failed\")\n",
    "    display(diff)\n",
    "# B,F\n",
    "diff = BF_t.diff(t).subs(t, tn).subs(om, c * knorm).expand() - dBF_dt\n",
    "diff.simplify()\n",
    "if diff != zeros(DD, 1):\n",
    "    print(\"Integration in time failed\")\n",
    "    display(diff)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Coefficients of PSATD equations in PML:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Code generation\n",
    "\n",
    "#       0,   1,   2,   3,   4,   5,   6,   7,   8,  9, 10, 11\n",
    "# EG: Exx, Exy, Exz, Eyx, Eyy, Eyz, Ezx, Ezy, Ezz, Gx, Gy, Gz\n",
    "# BF: Bxx, Bxy, Bxz, Byx, Byy, Byz, Bzx, Bzy, Bzz, Fx, Fy, Fz\n",
    "\n",
    "# Select update equation (left hand side)\n",
    "X_new = BF_new[0]\n",
    "\n",
    "# Extract individual terms (right hand side)\n",
    "for i in range(DD):\n",
    "    X = EG[i]\n",
    "    C1 = X_new.coeff(X, 1).simplify()\n",
    "    print(r\"Coefficient multiplying \" + str(X))\n",
    "    display(C1)\n",
    "    # print(ccode(Assignment(sp.symbols(r'LHS'), C1)))\n",
    "for i in range(DD):\n",
    "    X = BF[i]\n",
    "    C2 = X_new.coeff(X, 1).simplify()\n",
    "    print(r\"Coefficient multiplying \" + str(X))\n",
    "    display(C2)\n",
    "    # print(ccode(Assignment(sp.symbols(r'LHS'), C2)))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
